from typing import  Literal
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.messages import HumanMessage
# 导入状态图和状态
from langgraph.graph import END,StateGraph,MessagesState
#导入工具节点
from langgraph.prebuilt import ToolNode
from langchain_core.tools import tool
import os

#定义一个工具函数，用于代理调用外部工具
@tool
def search(query:str):
    """如果查询中包含“上海”，则返回“现在30度，有雾."""
    if "上海" in query.lower() or "Shanghai" in query.lower():
        return "现在30度，有雾"
    return "现在20度，有雨"

# 将工具函数放入到工具列表中
tools = [search]

#创建工具节点
tool_node = ToolNode(tools)

#初始化模型和工具，定义并绑定工具到模型中
model = (ChatOpenAI(model_name="deepseek-chat",
                   api_key=os.environ["DEEPSEEK_API_KEY"],
                   base_url=os.environ["DEEPSEEK_BASE_URL"])
         .bind_tools(tools))

# 定义路由函数，决定是否继续执行
def should_continue(state: MessagesState) -> Literal["tools",END]:
    # 获取最后一条消息
    messages = state["messages"]
    last_message = messages[-1]

    #如果LLM调用了工具，则转到“tools”节点
    if last_message.tool_calls :
        return "tools"
    #否则，停止执行
    return END

# 定义调用模型的函数
def call_model(state: MessagesState):
    # 获取最后一条消息
    messages = state["messages"]
    print("Messages before invoke:", messages)  # 调试输出
    response = model.invoke(messages)
    # 返回列表，因为这将被添加到现有列表中
    return {"messages": [response]}

# 用状态初始化图，定义一个新的状态图
workflow = StateGraph(MessagesState)

# 定义图节点，我们将循环的两个节点
workflow.add_node("agent", call_model)
workflow.add_node("tools",tool_node)


# 定义入口点和图边
# 设置入口点为“agent”
# 这意味着，这是第一个被调用的节点
workflow.set_entry_point("agent")


# 添加条件边
workflow.add_conditional_edges(
    # 首先 定义起始节点，我们使用agent
    # 这就意味着这些边式在调用 agent节点之后执行
    "agent",
    # 接下来传递决定下一个调用节点的函数
    should_continue,
)

# 添加从 tools 到 agent的普通边
# 这就意味着，调用 tools 之后，agent 将被调用
workflow.add_edge("tools", "agent")

#初始化内存，我们将使用MemorySaver节点来保存消息  可以存到redis、mongodb等
check_pointer = MemorySaver()

#将这个图，编译成可运行的langchain对象

app = workflow.compile(checkpointer = check_pointer)
final_state = app.invoke(
         {"messages": [HumanMessage(content="上海的天气如何")]},
         config={"configurable":{"thread_id":42}}
     )

print(final_state["messages"][-1].content)

final_state = app.invoke(
         {"messages": [HumanMessage(content="我问的是哪个城市？")]},
         config={"configurable":{"thread_id":42}}
     )

print(final_state["messages"][-1].content)

# 将生成的图片保存到文件中
graph_png = app.get_graph().draw_mermaid_png()
with open("code14-2-search.png", "wb") as f:
    f.write(graph_png)





